#质心计算过程
import numpy as np
import matplotlib.pyplot as plt
# 样本集
X = np.array([[1, 2], [2, 2], [6, 8], [7, 8]])
#定义初始化质心
C = np.array([[1.0, 2.0], [2.0, 2.0]])
m = len(X)
n_cluster = len(C)

plt.figure(figsize=[15, 6])
spr = 2
spc = 5
spn = 0
cmap = plt.cm.get_cmap('rainbow', n_cluster)


def xvisualize(xtitle):
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title(xtitle)
    plt.scatter(X[:, 0], X[:, 1], s=5, c=min_dis_idx, cmap=cmap)
    plt.scatter(C[:, 0], C[:, 1], marker='x', s=20, c=range(n_cluster), cmap=cmap)


#重复计算质心5次
iters = 0
while (iters<5) :
    iters += 1

    # fist cycle
    dis_matrix = np.zeros([m, n_cluster])
    for i, c in enumerate(C):
        #计算每个点到质心的欧式距离
        dis = np.sqrt(((X - c)**2).sum(axis=1))
        dis_matrix[:, i] = dis
    #求样本点属于哪一个类别
    min_dis_idx = np.argmin(dis_matrix, axis=1)
    xvisualize(f'{iters}th - determine clusters')

    # 2nd cycle
    for i in range(n_cluster):
        #更换每个质心的位置
        clustered_X = X[min_dis_idx == i]
        C[i] = np.mean(clustered_X, axis=0)
    xvisualize(f'{iters}th - determine new centers')

#打印所有样本的所属的簇
print('My code follow teachers')
print(min_dis_idx)
print(C)

plt.show()
